package cn.doitedu.api;

import beans.UserAction;
import com.alibaba.fastjson.JSON;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.contrib.streaming.state.EmbeddedRocksDBStateBackend;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import java.util.ArrayList;
import java.util.List;

public class _15_OperatorState_Demo {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setStateBackend(new HashMapStateBackend());
        env.setStateBackend(new EmbeddedRocksDBStateBackend(true));



        env.setParallelism(1);
        // 开启checkpoint
        env.enableCheckpointing(5000);
        // 指定checkpoint生成的状态快照数据文件存放在哪里
        env.getCheckpointConfig().setCheckpointStorage("file:///d:/ckpt");


        DataStreamSource<String> stream = env.socketTextStream("doitedu", 8899);

        SingleOutputStreamOperator<UserAction> beans = stream.map(json -> JSON.parseObject(json, UserAction.class));

        beans.map(new MyMapper()).print();


        env.execute();

    }

}

class MyMapper extends RichMapFunction<UserAction,Integer> implements CheckpointedFunction{

    ListState<UserAction> listState;
    List<UserAction> actionList = new ArrayList<>();

    public MyMapper(){

        System.out.println("构造方法调用了");
    }



    @Override
    public void open(Configuration parameters) throws Exception {


        System.out.println("open 调用了");
    }

    @Override
    public Integer map(UserAction userAction) throws Exception {
        actionList.add(userAction);
        return actionList.size();
    }

    /**
     * 对 flink底层做状态快照时，会调用本方法
     */
    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {

        System.out.println("snapshotState 调用了");

        // 在系统要做快照之前，我们赶紧把自己内存中存储的数据，添加到状态中
        listState.addAll(actionList);
    }

    /**
     * job重启，flink状态恢复后，会调用本方法
     */
    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        System.out.println("initializeState 调用了  ");

        // 获取一个operator state
        listState = context.getOperatorStateStore().getListState(new ListStateDescriptor<UserAction>("user_actions", UserAction.class));


        // 从状态中将数据恢复到我们自己的arrayList中
        for (UserAction userAction : listState.get()) {
            actionList.add(userAction);
        }



    }
}


